from scapy.all import sniff
import logging


class RTMPCapture:
    def __init__(self):
        self.rtmp_url = None
        self.stream_key = None
        self.on_error = None

        # 配置日志
        logging.basicConfig(
            level=logging.INFO,
            format="%(asctime)s - %(levelname)s - %(message)s",
            handlers=[logging.FileHandler("rtmp_capture.log"), logging.StreamHandler()],
        )
        self.logger = logging.getLogger(__name__)

        # 回调函数
        self.on_rtmp_url = None
        self.on_stream_key = None

    def packet_callback(self, packet):
        try:
            if packet.haslayer("TCP") and packet.haslayer("Raw"):
                payload = packet["Raw"].load

                # 解析RTMP URL
                # try:
                #     # 检查 rtmp:// 和 tcUrl
                #     rtmp_url = None
                #     rtmp_url_start = payload.find(b'rtmp://')
                #     if rtmp_url_start != -1:
                #         rtmp_url_end = payload.find(b'\x00', rtmp_url_start)
                #         if rtmp_url_end != -1:
                #             rtmp_url = payload[rtmp_url_start:rtmp_url_end].decode('utf-8')

                #     # 检查 tcUrl 参数
                #     tcurl_start = payload.find(b'tcUrl\x00')
                #     if tcurl_start != -1:
                #         # 跳过 tcUrl\x00\x02 (AMF字符串标记)
                #         tcurl_start = payload.find(b'\x02', tcurl_start)
                #         if tcurl_start != -1:
                #             # 获取字符串长度（2字节）
                #             str_len = (payload[tcurl_start+1] << 8) | payload[tcurl_start+2]
                #             tcurl_start += 3
                #             if tcurl_start + str_len <= len(payload):
                #                 rtmp_url = payload[tcurl_start:tcurl_start+str_len].decode('utf-8')

                #     if rtmp_url and self.on_rtmp_url:
                #         self.logger.info(f"服务器: {rtmp_url}")
                #         self.on_rtmp_url(rtmp_url)

                # except Exception as e:
                #     self.logger.error(f"解析RTMP URL时出错: {str(e)}")

                # 检查是否包含 releaseStream 命令并提取推流码
                if b"releaseStream" in payload or b"connect" in payload:
                    if b"connect" in payload:
                        self.logger.info("检测到 connect")
                    if b"releaseStream" in payload:
                        self.logger.info("检测到 releaseStream")
                    try:
                        pos = 0
                        while True:
                            pos = payload.find(b"\x02", pos)
                            if pos == -1:
                                break

                            if pos + 3 < len(payload):
                                str_len = (payload[pos + 1] << 8) | payload[pos + 2]
                                if pos + 3 + str_len <= len(payload):
                                    try:
                                        string_data = payload[
                                            pos + 3 : pos + 3 + str_len
                                        ].decode("utf-8")
                                        # 检查是否是推流码（通常较长且包含特定格式）
                                        if (
                                            b"connect" in payload
                                            and "rtmp" in string_data
                                        ):
                                            self.rtmp_url = string_data
                                            self.logger.info(
                                                f"RTMP Server: {string_data}"
                                            )
                                            if self.on_rtmp_url:
                                                self.on_rtmp_url(self.rtmp_url)
                                            break
                                        if (
                                            b"releaseStream" in payload
                                            and len(string_data) > 20
                                        ):
                                            self.stream_key = string_data
                                            self.logger.info(f"推流码: {string_data}")
                                            if self.on_stream_key:
                                                self.on_stream_key(self.stream_key)
                                        # if len(string_data) > 20:
                                        #     if b'connect' in payload:
                                        #         self.rtmp_url = string_data
                                        #         self.logger.info(f"RTMP Server: {string_data}")
                                        #         if self.on_rtmp_url:
                                        #             self.on_rtmp_url(self.rtmp_url)
                                        #         break
                                        #     if b'releaseStream' in payload:
                                        #         self.stream_key = string_data
                                        #         self.logger.info(f"推流码: {string_data}")
                                        #         if self.on_stream_key:
                                        #             self.on_stream_key(self.stream_key)
                                        #         break
                                    except Exception as e:
                                        self.logger.error(
                                            f"解析releaseStream or connect时出错: {str(e)}"
                                        )
                            pos += 1
                    except Exception as e:
                        self.logger.error(
                            f"处理releaseStream or connect命令时出错: {str(e)}"
                        )

                # # 检查是否包含其他可能的推流码格式
                # try:
                #     pos = 0
                #     while True:
                #         pos = payload.find(b'\x02', pos)
                #         if pos == -1:
                #             break

                #         if pos + 3 < len(payload):
                #             str_len = (payload[pos+1] << 8) | payload[pos+2]
                #             if pos + 3 + str_len <= len(payload):
                #                 try:
                #                     string_data = payload[pos+3:pos+3+str_len].decode('utf-8')
                #                     # 检查是否包含签名或其他标识符
                #                     if len(string_data) > 30 and ('sign=' in string_data or 'token=' in string_data):
                #                         self.stream_key = string_data
                #                         self.logger.info(f"推流码: {string_data}")
                #                         if self.on_stream_key:
                #                             self.on_stream_key(self.stream_key)
                #                         break
                #                 except:
                #                     pass
                #         pos += 1
                # except Exception as e:
                #     self.logger.error(f"检查其他推流码格式时出错: {str(e)}")

        except Exception as e:
            self.logger.error(f"处理数据包时出错: {str(e)}")

    def start_capture(self, filter_str="tcp"):
        """
        开始抓包
        :param filter_str: BPF 过滤器字符串
        """
        self.logger.info("开始抓包...")
        try:
            sniff(filter=filter_str, prn=self.packet_callback, store=0)
        except Exception as e:
            if self.on_error:
                self.on_error(str(e))
            self.logger.error(f"抓包过程出错: {str(e)}")


if __name__ == "__main__":
    capture = RTMPCapture()
    capture.start_capture(filter_str="tcp port 1935 or tcp port 80 or tcp port 443")
